#!/usr/bin/env python
# -*- coding:UTF-8

"""
Created on 2025-09-03

@author: powell
"""

import argparse
import logging
import logging.handlers
import os
import pwd
import re
import resource
import shutil
import subprocess
import sys
import time
import traceback
from textwrap import dedent

g_version = "v1.0"
g_tag_line = """# ====== Add by cspgsetup ======"""

# 这个程序执行时，logging的日志默认保存路径
g_log_file = "mongodb_setup.log"
logging.basicConfig(
    level=logging.DEBUG,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[
        logging.FileHandler(os.path.join('./', g_log_file)),
        logging.StreamHandler()
    ]
)


def get_disk_space_mb(path):
    """获取指定路径的磁盘空间信息（单位：MB）"""
    try:
        usage = shutil.disk_usage(path)
        MB = 1024 * 1024  # 1MB = 1024 * 1024字节

        return {
            "total": round(usage.total / MB, 2),
            "free": round(usage.free / MB, 2),
            "used": round(usage.used / MB, 2),
            "percent_used": round((usage.used / usage.total) * 100, 2)
            if usage.total > 0
            else 0,
        }
    except Exception as e:
        print(f"获取磁盘空间信息失败: {e}")
        return None


def run_cmd(cmd):
    logging.debug("Run: %s" % cmd)
    os.system(cmd)


def run_cmd_result(cmd):
    try:
        p = subprocess.Popen(
            cmd,
            shell=True,
            close_fds=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        out_msg = p.stdout.read()
        err_msg = p.stderr.read()
        err_code = p.wait()
        logging.debug("Run: %s, return %s" % (cmd, out_msg))
    except Exception as e:
        raise e

    return err_code, err_msg, out_msg


def get_target_file_list(dir_path, filter_str=None, file_type=None):
    """查询指定目录下符合条件的文件

    Args:
        dir_path (_type_): _description_
        filter_str (_type_, optional): _description_. Defaults to None.
        file_type (_type_, optional): _description_. Defaults to None.

    Returns:
        _type_: _description_
    """
    file_list = []
    for root, dirs, files in os.walk(dir_path):
        for file in files:
            is_target = False
            if filter_str and filter_str in file:
                is_target = True

            if file_type:
                if file.endswith(file_type):
                    is_target = True
                else:
                    is_target = False

            if is_target:
                file_list.append(os.path.join(root, file))

    return file_list


def add_os_user(os_uid, os_user):
    logging.info("Begin add user and group ...")
    cmd = "groupadd -g %d %s" % (os_uid, os_user)
    run_cmd(cmd)
    cmd = "useradd -m -u %(uid)d -g %(uid)d -s /bin/bash %(user)s" % {
        "uid": os_uid,
        "user": os_user,
    }
    run_cmd(cmd)
    logging.info("Add user and group ok.")

def get_memory_stat():
    mem = {}
    f = open("/proc/meminfo")
    lines = f.readlines()
    f.close()
    for line in lines:
        if len(line) < 2:
            continue
        name = line.split(":")[0]
        var = line.split(":")[1].split()[0]
        mem[name] = int(var) * 1024
    return mem

def modify_config_type1(config_file, modify_item_dict, deli_type=1, is_backup=True):
    """
    修改以等号或空格分隔的配置文件，默认为等号分隔的文件即（k = v）格式的配置文件，如postgresql.conf, /etc/sysctl.conf这些文件都是这种格式。
    :param config_file:
    :param modify_item_dict:
    :param deli_type: 1表示以等号分隔的配置文件，2表示以空格分隔的配置文件
    :param is_backup:
    :return:

    修改配置文件:
      1. 如果在文件中只有相应被注释掉的配置项，则新的配置项加在注释掉的配置项后面。
      2. 如果已存在配置项，则替换原有的配置项。
      3. 如果文件中不存在的配置项，则添加到文件尾
    例如modify_item_dict={'port':'5444'}，配置文件中只存在port的注释掉的配置：
      ...
      listen_addresses = '*'
      #port = 5432                            # (change requires restart)
      max_connections = 100                   # (change requires restart)
      ...

    执行后的结果是在此被注释掉的配置项后面加上新的配置项，结果变为：
      ...
      listen_addresses = '*'
      #port = 5432                            # (change requires restart)
      port = 5444
      max_connections = 100                   # (change requires restart)
      ...

    如果配置文件中存在port的注释掉的配置项和未被注释掉的相应的配置项：
      ...
      listen_addresses = '*'
      #port = 5432                            # (change requires restart)
      port = 5433
      max_connections = 100                   # (change requires restart)
      ...

    执行后的结果是在此被注释掉的配置项后面加上新的配置项，结果变为：
      ...
      listen_addresses = '*'
      #port = 5432                            # (change requires restart)
      port = 5444
      max_connections = 100                   # (change requires restart)
      ...

    :param config_file:
    :param modify_item_dict:
    :return:
    """

    fp = open(config_file, "r")
    ori_lines = fp.readlines()
    fp.close()

    # 下面的操作先找各个配置项的位置
    # item_line_num_dict1和item_line_num_dict2分别记录相应的配置项在文件中的行号。
    #   只是item_line_num_dict1字典中key是行号，而value是相应的配置项名称
    #   而item_line_num_dict2字典中key是配置项名称，而value是相应的行号
    item_line_num_dict1 = {}
    item_line_num_dict2 = {}

    # item_comment_line_num_dict1和item_comment_line_num_dict2分别记录配置文件中被注释掉的配置项在文件中的行号。
    item_comment_line_num_dict1 = {}
    item_comment_line_num_dict2 = {}

    i = 0
    for line in ori_lines:
        line = line.strip()
        if deli_type == 1:
            cells = line.split("=")
        else:
            cells = line.split()
        if len(cells) < 2:
            i += 1
            continue
        item_name = cells[0].strip()
        if item_name[0] == "#":
            if item_name[1:].strip() in modify_item_dict:
                item_comment_line_num_dict1[i] = item_name[1:]
                item_comment_line_num_dict2[item_name[1:]] = i
        if item_name in modify_item_dict:
            item_line_num_dict1[i] = item_name
            item_line_num_dict2[item_name] = i
        i += 1

    # 如果已存在相应的配置项，即使也存在注释掉的配置项，则就不能在已注释掉的配置项后再加上新配置项了，需要替换掉的已存在的配置项
    for item_name in item_comment_line_num_dict2:
        if item_name in item_line_num_dict2:
            i = item_comment_line_num_dict2[item_name]
            del item_comment_line_num_dict1[i]

    # 如果配置项在item_line_num_dict1中存在或在item_comment_line_num_dict1，则添加新配置项
    i = 0
    new_lines = []
    for line in ori_lines:
        line = line.strip()
        if i in item_line_num_dict1:
            if deli_type == 1:
                new_line = "%s = %s" % (
                    item_line_num_dict1[i],
                    modify_item_dict[item_line_num_dict1[i]],
                )
            else:
                new_line = "%s %s" % (
                    item_line_num_dict1[i],
                    modify_item_dict[item_line_num_dict1[i]],
                )
            new_lines.append(new_line)
        elif i in item_comment_line_num_dict1:
            # 如新行加到注释行的下一行处
            new_lines.append(line)
            item_name = item_comment_line_num_dict1[i]
            if deli_type == 1:
                new_line = "%s = %s" % (item_name, modify_item_dict[item_name])
            else:
                new_line = "%s %s" % (item_name, modify_item_dict[item_name])
            new_lines.append(new_line)
        else:
            new_lines.append(line)
        i += 1

    # 把配置文件中不存在的配置项，添加到文件尾（按item_name排序）
    items = sorted(modify_item_dict.keys())
    for item_name in items:
        if (
            item_name not in item_line_num_dict2
            and item_name not in item_comment_line_num_dict2
        ):
            if deli_type == 1:
                new_line = "%s = %s" % (item_name, modify_item_dict[item_name])
            else:
                new_line = "%s %s" % (item_name, modify_item_dict[item_name])
            new_lines.append(new_line)

    if is_backup:
        config_file_bak = "%s.%s" % (config_file, time.strftime("%Y%m%d%H%M%S"))
        shutil.copy(config_file, config_file_bak)

    fp = open(config_file, "w")
    content = "\n".join(new_lines)
    fp.write(content)
    fp.close()


def sshd_config_set():
    # 修改sshd_config中的配置
    # /etc/ssh/sshd_config
    ci_dict = {"UseDNS": "no", "GSSAPIAuthentication": "no"}
    modify_config_type1("/etc/ssh/sshd_config", ci_dict, deli_type=2)


def init_os_env():
    # 修改sshd_config
    sshd_config_set()

    # 关闭防火墙
    run_cmd("systemctl disable firewalld")

    # io的默认调度策略改成deadline
    run_cmd('''grubby --update-kernel=ALL --args="elevator=deadline"''')

    # 禁用透明大页
    run_cmd('''grubby --update-kernel=ALL --args="transparent_hugepage=never"''')


def init_options_parser():
    """
    :return:
    """

    usage = f"%prog {sys.argv[1]} [options]\n"
    parser = argparse.ArgumentParser(usage=usage)
    parser.add_argument(
        "-l",
        "--loglevel",
        action="store",
        dest="loglevel",
        default="info",
        help="Specifies log level:  debug, info, warn, error, critical, default is info",
    )
    return parser

def action_init_os_env():
    parser = init_options_parser()
    try:
        _options = parser.parse_args(sys.argv[1:])
    except Exception:
        print(traceback.format_exc())
    init_os_env()

def install_deps():
    """
    安装依赖包
    """
    cmd = "yum install cyrus-sasl cyrus-sasl-gssapi cyrus-sasl-plain krb5-libs libcurl openldap openssl xz-libs"
    run_cmd(cmd)

def install_mongo_server():
    """
    安装MongoDB Server
    """
    cmd = "rpm -ivh /root/mongodb-enterprise-server-8.0.13-1.el8.x86_64.rpm "
    run_cmd(cmd)


def install_mongo_shell(mongosh_tgz_file = "mongosh-2.5.7-linux-x64.tgz"):
    """
    安装MongoDB Shell
    """
    mongosh_dir = '/usr/lib/mongosh'
    cmd = f"install -d {mongosh_dir} && tar xvf /root/{mongosh_tgz_file}  -C {mongosh_dir}"
    run_cmd(cmd)
    cmd = f"ls {mongosh_dir}"
    code, errmsg, outmsg = run_cmd_result(cmd)
    if code == 0:
        print("MongoDB Shell installed successfully.")
    else:
        print("Failed to install MongoDB Shell.")
        print("Error message: %s", errmsg)
        print("Output message: %s", outmsg)
    # 创建软链接
    str_outmsg = str(outmsg, 'utf-8').strip()
    cmd = f"ln -s {mongosh_dir}/{str_outmsg}/bin/mongosh /usr/bin/mongosh"
    run_cmd(cmd)
    # 检查是否安装成功
    cmd = f"mongosh --version"
    code, errmsg, outmsg = run_cmd_result(cmd)
    if code == 0:
        print("MongoDB Shell installed successfully.")
    else:
        print("Failed to install MongoDB Shell.")
        print("Error message: %s", errmsg)
        print("Output message: %s", outmsg)

def main():
    # prepare,先告诉用户，需要做的准备工作
    # 下载mongo server端rpm包，mongodb-enterprise-server-8.0.13-1.el8.x86_64.rpm，以及mongosh-2.5.7-linux-x64.tgz 客户端,并放到/root/目录下
    # 打印上述信息给用户确认，用户确认后继续，否则退出
    #下面内容输出到屏幕
    print("Please confirm the following infos:")
    print(" - Please ensure your OS is el8.x86_64.")
    print(" - Please ensure you have internet access.")
    print(" - Please ensure the following files are in /root/:")
    print(" -- mongodb-enterprise-server-8.0.13-1.el8.x86_64.rpm")
    print(" -- mongosh-2.5.7-linux-x64.tgz")
    confirm = input("Are these files ready? (y/n): ")
    if confirm.lower() != "y":
        print("User aborted.")
        sys.exit(1)

    # 初始化操作系统环境
    init_os_env()
    # 安装依赖
    install_deps()
    # 安装MongoDB server
    install_mongo_server()

    # 启动 service
    run_cmd("systemctl start mongod.service")
    # 输出状态
    code, errmsg, outmsg = run_cmd_result("systemctl status mongod.service")
    if code == 0:
        print("MongoDB service is running.")
    else:
        print("MongoDB service is not running.")
        print("Error message: %s", errmsg)
        print("Output message: %s", outmsg)

    # 安装MongoDB Shell
    install_mongo_shell()

    sys.exit(1)


if __name__ == "__main__":
    main()

